# 导入json模块，用于处理JSON数据
import json

# 导入logging模块，用于日志记录
import logging

# 导入requests模块，用于发送HTTP请求
import requests

# 从urllib.parse模块导入urljoin，用于拼接URL
from urllib.parse import urljoin

# 配置日志级别为INFO
logging.basicConfig(level=logging.INFO)
# 获取当前模块的日志记录器
logger = logging.getLogger(__name__)


# 定义MCPClient类，表示MCP HTTP客户端
class MCPClient:
    """MCP HTTP客户端"""

    # 初始化方法，设置服务器URL和相关属性
    def __init__(self, server_url: str):
        # 协议版本号
        self.protocol_version = "2025-06-18"
        # 客户端信息
        self.client_info = {"name": "MCP客户端", "version": "1.0.0"}
        # 服务器信息，初始化为None
        self.server_info = None
        # 服务器能力，初始化为None
        self.capabilities = None
        # 会话ID，初始化为None
        self.session_id = None
        # 消息ID，初始值为1
        self.message_id = 1

        # 去除服务器URL末尾的斜杠
        self.server_url = server_url.rstrip("/")
        # 拼接MCP接口的完整URL
        self.mcp_endpoint = urljoin(self.server_url, "/mcp")
        # 创建requests的Session对象
        self.session = requests.Session()
        # 设置Session的默认请求头
        self.session.headers.update(
            {
                "Accept": "application/json",
                "MCP-Protocol-Version": self.protocol_version,
            }
        )

    def call_long_running_tool(self, tool_name, arguments):
        # 构造调用工具的请求消息
        request = {
            "jsonrpc": "2.0",
            "id": self._generate_id(),
            "method": "tools/call",
            "params": {"calls": [{"name": tool_name, "arguments": arguments}]},
        }

        # 发送请求并获取响应
        response = self._send_message(request)
        # 如果响应存在且包含result字段
        if response and "result" in response:
            # 获取调用结果
            calls = response["result"].get("calls", [])
            if calls:
                return calls[0]
        else:
            # 调用失败，记录错误
            logger.error(f"调用工具 {tool_name} 失败")
        return None

    def cancel_request(self, request_id, reason):
        request = {
            "jsonrpc": "2.0",
            "id": self._generate_id(),
            "method": "notifications/cancelled",
            "params": {"request_id": request_id, "reason": reason},
        }

        # 发送请求并获取响应
        response = self._send_message(request)
        # 如果响应存在且包含result字段
        if response and "result" in response:
            result = response["result"]
            status = result.get("status")
            if status == "cancelled":
                logger.info(f"请求{request_id}已经被取消")
                return True
            else:
                logger.warning(f"请求{request_id}取消失败 {status}")
                return False
        else:
            logger.warning(f"请求{request_id}取消失败 ")
            return False

    # 重置会话方法
    def reset_session(self):
        """重置会话"""
        # 保存旧的会话ID
        old_session_id = self.session_id
        # 清空当前会话ID
        self.session_id = None
        # 记录会话重置信息
        logger.info(f"会话已重置，原会话ID: {old_session_id}")

    # 获取当前会话信息方法
    def get_session_info(self):
        """获取当前会话信息"""
        # 返回会话相关信息字典
        return {
            "session_id": self.session_id,
            "server_url": self.server_url,
            "protocol_version": self.protocol_version,
        }

    # 获取可用工具列表方法
    def list_tools(self):
        """获取可用工具列表"""
        # 构造请求消息
        request = {
            "jsonrpc": "2.0",
            "id": self._generate_id(),
            "method": "tools/list",
            "params": {},
        }

        # 发送请求并获取响应
        response = self._send_message(request)
        # 如果响应存在且包含result字段
        if response and "result" in response:
            # 获取工具列表
            tools = response["result"].get("tools", [])
            logger.info(f"可用工具数量: {len(tools)}")
            return tools
        else:
            # 获取失败，记录错误
            logger.error("获取工具列表失败")
            return []

    # 通过HTTP发送消息的私有方法
    def _send_message(self, message):
        """通过HTTP发送消息"""
        try:
            # 初始化请求头字典
            headers = {}
            # 如果有会话ID，则添加到请求头
            if self.session_id:
                headers["Mcp-Session-Id"] = self.session_id
                logger.debug(f"使用会话ID: {self.session_id}")

            # 发送POST请求
            response = self.session.post(
                self.mcp_endpoint, json=message, headers=headers, timeout=30
            )

            # 检查响应状态码
            if response.status_code == 200:
                # 检查响应头中是否有新的会话ID
                new_session_id = response.headers.get("Mcp-Session-Id")
                # 如果有新的会话ID且与当前不同，则更新
                if new_session_id and new_session_id != self.session_id:
                    self.session_id = new_session_id
                    logger.info(f"更新会话ID: {self.session_id}")

                # 返回响应的JSON内容
                return response.json()
            # 如果状态码为202，表示已接受但无内容
            elif response.status_code == 202:
                # Accepted - 通知或响应
                return None
            # 如果状态码为400，表示请求错误
            elif response.status_code == 400:
                logger.error(f"请求错误: {response.text}")
                return None
            # 如果状态码为401，表示会话ID无效
            elif response.status_code == 401:
                logger.error("会话ID无效，需要重新初始化")
                self.session_id = None
                return None
            # 其他HTTP错误
            else:
                logger.error(f"HTTP错误: {response.status_code} - {response.text}")
                return None

        # 捕获请求异常
        except requests.exceptions.RequestException as e:
            logger.error(f"HTTP请求错误: {e}")
            return None

    # 生成消息ID的方法
    def _generate_id(self) -> str:
        """生成消息ID"""
        # 消息ID自增
        self.message_id += 1
        # 返回字符串类型的消息ID
        return str(self.message_id)

    # 发送消息到服务器的方法（重复定义，实际只会用后面这个）
    def _send_message(self, message):
        """通过HTTP发送消息"""
        try:
            # 初始化请求头字典
            headers = {}
            # 如果有会话ID，则添加到请求头
            if self.session_id:
                headers["Mcp-Session-Id"] = self.session_id
                logger.debug(f"使用会话ID: {self.session_id}")

            # 发送POST请求到MCP端点
            response = self.session.post(
                self.mcp_endpoint, json=message, headers=headers, timeout=30
            )

            # 检查响应状态码
            if response.status_code == 200:
                # 从响应头获取新的会话ID
                new_session_id = response.headers.get("Mcp-Session-Id")
                # 如果有新的会话ID且与当前不同，则更新
                if new_session_id and new_session_id != self.session_id:
                    self.session_id = new_session_id
                    logger.info(f"更新会话ID: {self.session_id}")

                # 返回响应的JSON内容
                return response.json()
            # 如果状态码为202，表示已接受但无内容
            elif response.status_code == 202:
                # Accepted - 通知或响应
                return None
            # 如果状态码为400，表示请求错误
            elif response.status_code == 400:
                logger.error(f"请求错误: {response.text}")
                return None
            # 如果状态码为401，表示会话ID无效
            elif response.status_code == 401:
                logger.error("会话ID无效，需要重新初始化")
                self.session_id = None
                return None
            # 其他HTTP错误
            else:
                logger.error(f"HTTP错误: {response.status_code} - {response.text}")
                return None

        # 捕获请求异常
        except requests.exceptions.RequestException as e:
            logger.error(f"HTTP请求错误: {e}")
            return None

    # 调用工具方法
    def call_tool(self, tool_name, arguments):
        """调用工具"""
        # 构造调用工具的请求消息
        request = {
            "jsonrpc": "2.0",
            "id": self._generate_id(),
            "method": "tools/call",
            "params": {"calls": [{"name": tool_name, "arguments": arguments}]},
        }

        # 发送请求并获取响应
        response = self._send_message(request)
        # 如果响应存在且包含result字段
        if response and "result" in response:
            # 获取调用结果
            calls = response["result"].get("calls", [])
            if calls:
                return calls[0]
        else:
            # 调用失败，记录错误
            logger.error(f"调用工具 {tool_name} 失败")
        return None

    # 初始化与服务器连接的方法
    def initialize(self) -> bool:
        """初始化与服务器的连接"""
        # 构造初始化请求消息
        request = {
            "jsonrpc": "2.0",
            "id": self._generate_id(),
            "method": "initialize",
            "params": {
                "protocolVersion": self.protocol_version,
                "capabilities": {},
                "clientInfo": self.client_info,
            },
        }

        # 发送初始化请求
        response = self._send_message(request)
        # 如果响应存在且包含result字段
        if response and "result" in response:
            result = response["result"]
            # 获取服务器信息
            self.server_info = result.get("serverInfo")
            # 获取服务器能力
            self.capabilities = result.get("capabilities")

            # 提取并存储会话ID
            if "sessionId" in result:
                self.session_id = result["sessionId"]
                logger.info(f"获取到会话ID: {self.session_id}")

            logger.info(f"服务器初始化成功: {self.server_info}")
            return True
        else:
            # 初始化失败，记录错误
            logger.error("服务器初始化失败")
            return False


# HTTP传输的函数
def http_transport():
    """HTTP传输"""
    print("=== HTTP传输 ===")

    client = MCPClient("http://localhost:8000")

    try:
        if not client.initialize():
            print("初始化失败")
            return
        result = client.call_long_running_tool(
            "long_running_task", {"duration": 20, "task_id": "long_task_001"}
        )
        if result:
            print(f"任务启动结果:{result}")
            import time

            time.sleep(10)
            print(f"向服务器发送取消任务的请求")
            if client.cancel_request("long_task_001", "我要演示取消任务"):
                print(f"取消任务成功")
            else:
                print(f"取消任务失败")

        else:
            print(f"任务启动失败")
    except Exception as e:
        print(f"HTTP错误: {e}")


# 主函数
def main():
    """主函数"""
    # 调用HTTP传输函数
    http_transport()


# 判断是否为主程序入口
if __name__ == "__main__":
    main()
